#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/des.h>
#include <openssl/evp.h>
#include "../include/crypto.h"

int
initDes( int mode, char *password, char *key, char *ivec, DES_key_schedule *ks, DES_cblock *iv )
{
	DES_cblock k;

	if( (mode != ECB) && (mode != CFB) && (mode != OFB) && (mode != CBC) )
		return 0;

	if( ( ks == NULL ) || (iv == NULL ) )
		return 0;

	if( (mode == ECB) && (password == NULL) )
		return 0;
	
	if( (mode == CFB) || (mode == OFB) || (mode == CBC) )
		if( (password == NULL) && ((key == NULL) && (ivec == NULL)) )
			return 0;

	if( password != NULL ){
		/* Extracts k and iv from password */	
		const EVP_CIPHER *cipher;
		const EVP_MD *dgst = EVP_md5(); 

		switch( mode ){
			case ECB: 
				cipher = EVP_des_ecb();
				break;
			case CFB: 
				cipher = EVP_des_cfb8();
				break;
			case OFB: 
				cipher = EVP_des_ofb();
				break;
			case CBC: 
				cipher = EVP_des_cbc();
				break;
			default: break;
		}

		EVP_BytesToKey( cipher, dgst, NULL, password, strlen(password), 1, k, *iv );
	}
	else{
		memcpy( k, key, DES_KEY_SIZE );
		memcpy( *iv, ivec, DES_KEY_SIZE );
	}
	
	/* Sets the parity of the passed key to odd */
	DES_set_odd_parity( &k );
	
	/* Generates the key schedule */
	DES_set_key_checked( &k, ks );

	return 1;
}


char *
des( int operation, int mode, DES_key_schedule *ks, DES_cblock *iv, char *block )
{
	char *dblock;
	
	if( ((operation != DES_ENCRYPT) && (operation != DES_DECRYPT)) ||
        ((mode != ECB) && (mode != CFB) && (mode != OFB) && (mode != CBC)) ||
        (block == NULL) )
		return NULL;

	if( ks == NULL )
		return NULL;

	if( (mode == CFB) || (mode == OFB) || (mode == CBC) )
		if( iv == NULL )
			return NULL;

	if( ( dblock = (char *) calloc(DES_BLOCK_SIZE,1) ) == NULL )
		return NULL;
	
	switch( mode ){
		case ECB:
			DES_ecb_encrypt( (const_DES_cblock *) block, (DES_cblock *) dblock, ks, operation );
			break;
		case CFB:
			DES_cfb_encrypt( (const unsigned char *) block, (unsigned char *) dblock, DES_BLOCK_SIZE, DES_BLOCK_SIZE, ks, iv, operation );
			break;
		case OFB:
			DES_ofb_encrypt( (const unsigned char *) block, (unsigned char *) dblock, DES_BLOCK_SIZE, DES_BLOCK_SIZE, ks, iv );
			break;
		case CBC:
			DES_ncbc_encrypt( (const unsigned char *) block, (unsigned char *) dblock, DES_BLOCK_SIZE, ks, iv, operation );
			break;
		default: break;
	}

	return dblock;
}

int
initAes( int operation, int mode, char *password, char *key, char *ivec, AES_KEY *aesKey, unsigned char *iv )
{
	unsigned char k[AES_KEY_SIZE];
	
	if( ((operation != AES_ENCRYPT) && (operation != AES_DECRYPT)) || ((mode != ECB) && (mode != CFB) && (mode != OFB) && (mode != CBC)) )
		return 0;

	if( (aesKey == NULL) || (iv == NULL) )
		return 0;

	if( (mode == ECB) && (password == NULL) )
		return 0;
	
	if( (mode == CFB) || (mode == OFB) || (mode == CBC) )
		if( (password == NULL) && ((key == NULL) && (ivec == NULL)) )
			return 0;

	if( password != NULL ){
		/* Extracts k and iv from password */	
		const EVP_CIPHER *cipher;
		const EVP_MD *dgst = EVP_md5();

		switch( mode ){
			case ECB:
				cipher = EVP_aes_128_ecb();
				break;
			case CFB:
				cipher = EVP_aes_128_cfb();
				break;
			case OFB: 
				cipher = EVP_aes_128_ofb();
				break;
			case CBC: 
				cipher = EVP_aes_128_cbc();
				break;
			default:
				break;
		}

		EVP_BytesToKey( cipher, dgst, NULL, password, strlen(password), 1, k, iv );
	}
	else{
		memcpy( k, key, AES_KEY_SIZE );
		memcpy( iv, ivec, AES_KEY_SIZE );
	}
	
	if( operation == AES_ENCRYPT )
		AES_set_encrypt_key( k, 128, aesKey );
	else
		AES_set_decrypt_key( k, 128, aesKey );

	return 1;
}

char *
aes( int operation, int mode, AES_KEY *aesKey, unsigned char *iv, char *block )
{
	char *ablock;
	
	if( ((operation != AES_ENCRYPT) && (operation != AES_DECRYPT)) || ((mode != ECB) && (mode != CFB) && (mode != OFB) && (mode != CBC)) || (block == NULL) )
		return NULL;

	if( aesKey == NULL )
		return NULL;
	
	if( (mode == CFB) || (mode == OFB) || (mode == CBC) )
		if( iv == NULL )
			return NULL;

	if( ( ablock = (char *) calloc(AES_BLOCK_SIZE,1) ) == NULL )
		return NULL;
	
	int num = AES_BLOCK_SIZE;

	switch( mode ){
		case ECB:
			AES_ecb_encrypt( (const unsigned char *) block, (unsigned char *) ablock, aesKey, operation );
			break;
		case CFB:
			AES_cfb128_encrypt( (const unsigned char *) block, (unsigned char *) ablock, AES_BLOCK_SIZE, aesKey, iv, &num, operation );
			break;
		case OFB:
			AES_ofb128_encrypt( (const unsigned char *) block, (unsigned char *) ablock, AES_BLOCK_SIZE, aesKey, iv, &num );
			break;
		case CBC:
			AES_cbc_encrypt( (const unsigned char *) block, (unsigned char *) ablock, AES_BLOCK_SIZE, aesKey, iv, operation );
			break;
		default:
			break;
	}
	
	return ablock;
}


